// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions::test;

namespace {

class ArrayMinMaxByTest : public FunctionBaseTest {
 protected:
  template <typename T>
  void testArrayMinMaxBy(
      const std::vector<std::optional<T>>& expected,
      const std::string& expression,
      const RowVectorPtr inputVector) {
    auto result = evaluate(expression, inputVector);
    kumo::pollux::test::assertEqualVectors(
        make_nullable_flat_vector<T>(expected), result);
  }

  template <typename T>
  void testArrayMinMaxByComplexType(
      ArrayVectorPtr input,
      const std::string& expression,
      T expected) {
    auto inputRow = make_row_vector({input});
    auto result = evaluate(expression, inputRow);
    kumo::pollux::test::assertEqualVectors(expected, result);
  }
};

TEST_F(ArrayMinMaxByTest, maxByBasic) {
  // array of integers
  auto intInput = make_row_vector({make_nullable_array_vector<int32_t>({
      {1, 2},
      {-3, 2},
      {-3, 3, 0},
      {3, 3, 3},
      {7},
      {7, std::nullopt},
      {std::nullopt},
  })});

  std::vector<std::optional<int32_t>> intExpected1{
      2,
      2,
      3,
      3,
      7,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<int32_t>(
      intExpected1, "array_max_by(c0, i -> (i))", intInput);
  std::vector<std::optional<int32_t>> intExpected2{
      1,
      -3,
      -3,
      3,
      7,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<int32_t>(
      intExpected2, "array_max_by(c0, i -> ((-1) * i))", intInput);

  // array of doubles
  auto doubleInput = make_row_vector({make_nullable_array_vector<double>({
      {1.0, 2.0},
      {-3.0, 2.0},
      {-3.0, 3.0, 0.0},
      {7.0},
      {7.0, std::nullopt},
      {std::nullopt},
  })});
  std::vector<std::optional<double>> doubleExpected1{
      2.0,
      2.0,
      3.0,
      7.0,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<double>(
      doubleExpected1, "array_max_by(c0, i -> (i))", doubleInput);
  std::vector<std::optional<double>> doubleExpected2{
      2.0,
      -3.0,
      3.0,
      7.0,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<double>(
      doubleExpected2, "array_max_by(c0, i -> (i * i))", doubleInput);
  std::vector<std::optional<double>> doubleExpected3{
      1.0,
      -3.0,
      -3.0,
      7.0,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<double>(
      doubleExpected3, "array_max_by(c0, i -> (1.0 - i))", doubleInput);

  // array of strings
  auto stringInput = make_row_vector({make_nullable_array_vector<StringView>({
      {"1", "2"},
      {"a", "bb", "c"},
      {"abc", std::nullopt},
      {std::nullopt},
  })});
  std::vector<std::optional<StringView>> stringExpected1{
      "2",
      "bb",
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<StringView>(
      stringExpected1, "array_max_by(c0, i -> LENGTH(i))", stringInput);
}

TEST_F(ArrayMinMaxByTest, maxByComplexTypes) {
  // array of arrays
  auto intArrayInput = make_nested_array_vector_from_json<int64_t>(
      {"[[1, 2, 3]]",
       "[[1, 2, 3], [3, 3], [4, 4], [5, 5]]",
       "[[1, 1], [3, 3], [4, 4], [5, 5]]",
       "[[1, 1], [3, 3], [4, 4], [5, 5], [null]]",
       "[[1, 1], [3, 3], [null, null, null], [4, 4], [5, 5]]",
       "[[], []]"});
  auto intArrayResult = make_array_vector_from_json<int64_t>(
      {"[1, 2, 3]",
       "[1, 2, 3]",
       "[5, 5]",
       "[5, 5]",
       "[null, null, null]",
       "[]"});
  testArrayMinMaxByComplexType<ArrayVectorPtr>(
      intArrayInput, "array_max_by(c0, i -> CARDINALITY(i))", intArrayResult);

  // array of maps
  std::vector<std::pair<StringView, std::optional<int64_t>>> a{
      {"one"_sv, 1}, {"two"_sv, 2}};
  std::vector<std::pair<StringView, std::optional<int64_t>>> b{{"one"_sv, 0}};
  std::vector<std::pair<StringView, std::optional<int64_t>>> c{
      {"three"_sv, 4}, {"four"_sv, 5}};
  std::vector<
      std::vector<std::vector<std::pair<StringView, std::optional<int64_t>>>>>
      mapData{{a, b}, {b}, {c}, {c, a}, {a, c}};
  auto array_vector = make_array_of_map_vector<StringView, int64_t>(mapData);
  auto mapResult = make_map_vector_from_json<std::string, int64_t>({
      "{\"one\": 1, \"two\": 2}",
      "{\"one\": 0}",
      "null",
      "null",
      "null",
  });
  testArrayMinMaxByComplexType<MapVectorPtr>(
      array_vector, "array_max_by(c0,x -> x['one'])", mapResult);

  RowTypePtr rowType = ROW({INTEGER(), VARCHAR()});

  using ArrayOfRow = std::vector<std::optional<std::tuple<int, std::string>>>;
  std::vector<ArrayOfRow> rowData = {
      {{{1, "one"}}, {{2, "two"}}, {{3, "three"}}},
      {{{1, "one"}}, {{4, "three"}}, {{3, "two"}}},
      {{{1, "one"}}}};
  auto ArrayOfRows = make_array_of_row_vector(rowData, rowType);

  auto c0 = make_nullable_flat_vector<int32_t>({3, 4, 1});
  auto c1 = make_nullable_flat_vector<std::string>({"three", "three", "one"});
  auto resultVector = make_row_vector({c0, c1});

  testArrayMinMaxByComplexType<RowVectorPtr>(
      ArrayOfRows, "array_max_by(c0,r -> r)", resultVector);
} // namespace

TEST_F(ArrayMinMaxByTest, minByBasic) {
  // array of integers
  auto intInput = make_row_vector({make_nullable_array_vector<int32_t>({
      {1, 2},
      {-3, 2},
      {-3, 3, 0},
      {3, 3, 3},
      {7},
      {7, std::nullopt},
      {std::nullopt},
  })});
  std::vector<std::optional<int32_t>> intExpected1{
      1,
      -3,
      -3,
      3,
      7,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<int32_t>(
      intExpected1, "array_min_by(c0, i -> (i))", intInput);
  std::vector<std::optional<int32_t>> intExpected2{
      2,
      2,
      3,
      3,
      7,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<int32_t>(
      intExpected2, "array_min_by(c0, i -> ((-1) * i))", intInput);

  // array of doubles
  auto doubleInput = make_row_vector({make_nullable_array_vector<double>({
      {1.0, 2.0},
      {-3.0, 2.0},
      {-3.0, 3.0, 0.0},
      {7.0},
      {7.0, std::nullopt},
      {std::nullopt},
  })});
  std::vector<std::optional<double>> doubleExpected1{
      1.0,
      -3.0,
      -3.0,
      7.0,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<double>(
      doubleExpected1, "array_min_by(c0, i -> (i))", doubleInput);
  std::vector<std::optional<double>> doubleExpected2{
      1.0,
      2.0,
      0.0,
      7.0,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<double>(
      doubleExpected2, "array_min_by(c0, i -> (i * i))", doubleInput);
  std::vector<std::optional<double>> doubleExpected3{
      2.0,
      2.0,
      3.0,
      7.0,
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<double>(
      doubleExpected3, "array_min_by(c0, i -> (1.0 - i))", doubleInput);

  // array of strings
  auto stringInput = make_row_vector({make_nullable_array_vector<StringView>({
      {"1", "2"},
      {"a", "bb", "c"},
      {"abc", std::nullopt},
      {std::nullopt},
  })});
  std::vector<std::optional<StringView>> stringExpected1{
      "2",
      "c",
      std::nullopt,
      std::nullopt,
  };
  testArrayMinMaxBy<StringView>(
      stringExpected1, "array_min_by(c0, i -> LENGTH(i))", stringInput);
}

TEST_F(ArrayMinMaxByTest, minByComplexTypes) {
  // array of arrays
  auto intArrayInput = make_nested_array_vector_from_json<int64_t>(
      {"[[1, 2, 3]]",
       "[[1, 2, 3], [3, 3], [4, 4], [5, 5]]",
       "[[1, 1], [3, 3], [4, 4], [5, 5]]",
       "[[1, 1], [3, 3], [4, 4], [5, 5], [null]]",
       "[[1, 1], [3, 3, 3], [null, null, null], [4, 4, 4], [5, 5, 5]]",
       "[[], []]"});
  auto intArrayResult = make_array_vector_from_json<int64_t>(
      {"[1, 2, 3]", "[5, 5]", "[5, 5]", "[null]", "[1, 1]", "[]"});
  testArrayMinMaxByComplexType<ArrayVectorPtr>(
      intArrayInput, "array_min_by(c0, i -> CARDINALITY(i))", intArrayResult);

  // array of maps
  std::vector<std::pair<StringView, std::optional<int64_t>>> a{
      {"one"_sv, 1}, {"two"_sv, 2}};
  std::vector<std::pair<StringView, std::optional<int64_t>>> b{{"one"_sv, 0}};
  std::vector<std::pair<StringView, std::optional<int64_t>>> c{
      {"three"_sv, 4}, {"four"_sv, 5}};
  std::vector<
      std::vector<std::vector<std::pair<StringView, std::optional<int64_t>>>>>
      mapData{{a, b}, {b}, {c}, {c, a}, {a, c}};
  auto array_vector = make_array_of_map_vector<StringView, int64_t>(mapData);
  auto mapResult = make_map_vector_from_json<std::string, int64_t>({
      "{\"one\": 0}",
      "{\"one\": 0}",
      "null",
      "null",
      "null",
  });
  testArrayMinMaxByComplexType<MapVectorPtr>(
      array_vector, "array_min_by(c0,x -> x['one'])", mapResult);

  RowTypePtr rowType = ROW({INTEGER(), VARCHAR()});

  using ArrayOfRow = std::vector<std::optional<std::tuple<int, std::string>>>;
  std::vector<ArrayOfRow> rowData = {
      {{{1, "one"}}, {{2, "two"}}, {{3, "three"}}},
      {{{1, "one"}}, {{4, "three"}}, {{3, "two"}}},
      {{{1, "one"}}}};
  auto ArrayOfRows = make_array_of_row_vector(rowData, rowType);

  auto c0 = make_nullable_flat_vector<int32_t>({1, 1, 1});
  auto c1 = make_nullable_flat_vector<std::string>({"one", "one", "one"});
  auto resultVector = make_row_vector({c0, c1});

  testArrayMinMaxByComplexType<RowVectorPtr>(
      ArrayOfRows, "array_min_by(c0, r -> r)", resultVector);
} // namespace
} // namespace
